import cv2
import numpy as np

# 读取图像
image = cv2.imread("img_1.png")  # 替换为你的图像路径
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # 转为灰度图

# 二值化（调整阈值，确保灯条清晰显示为白色）
_, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
# 查找所有轮廓（RETR_LIST：获取所有轮廓，CHAIN_APPROX_NONE：保留所有轮廓点）
contours, _ = cv2.findContours(binary, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)

valid_contours = []  # 存储筛选后有效的灯条轮廓

for contour in contours:
    # 条件1：过滤面积过小/过大的轮廓（根据实际灯条大小调整范围）
    area = cv2.contourArea(contour)
    if area < 500 or area > 5000:
        continue

    # 条件2：多边形拟合（逼近为六边形，边数约为6）
    perimeter = cv2.arcLength(contour, True)  # 计算轮廓周长
    epsilon = 0.03 * perimeter  # 拟合精度（越小越接近原轮廓）
    approx = cv2.approxPolyDP(contour, epsilon, True)  # 多边形拟合

    # 筛选边数接近6的轮廓（箭头灯条抽象为六边形）
    if len(approx) >= 5 and len(approx) <= 7:
        valid_contours.append(contour)
result = image.copy()  # 复制原图用于绘制结果

for contour in valid_contours:
    # 计算最小包围矩形（旋转矩形，更贴合灯条形状）
    rect = cv2.minAreaRect(contour)
    box = cv2.boxPoints(rect)  # 获取矩形的四个顶点
    box = np.int32(box)  # 转为整数坐标

    # 绘制最小包围矩形（绿色，线宽2）
    cv2.drawContours(result, [box], 0, (0, 255, 0), 2)
cv2.imshow("Result", result)
cv2.waitKey(0)
cv2.destroyAllWindows()